{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# PyTorch Profiler With TensorBoard\nThis tutorial demonstrates how to use TensorBoard plugin with PyTorch Profiler\nto detect performance bottlenecks of the model.\n\n## Introduction\nPyTorch 1.8 includes an updated profiler API capable of\nrecording the CPU side operations as well as the CUDA kernel launches on the GPU side.\nThe profiler can visualize this information\nin TensorBoard Plugin and provide analysis of the performance bottlenecks.\n\nIn this tutorial, we will use a simple Resnet model to demonstrate how to\nuse TensorBoard plugin to analyze model performance.\n\n## Setup\nTo install ``torch`` and ``torchvision`` use the following command:\n\n::\n\n   pip install torch torchvision\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Steps\n\n1. Prepare the data and model\n2. Use profiler to record execution events\n3. Run the profiler\n4. Use TensorBoard to view results and analyze model performance\n5. Improve performance with the help of profiler\n6. Analyze performance with other advanced features\n\n### 1. Prepare the data and model\n\nFirst, import all necessary libraries:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nimport torch.nn\nimport torch.optim\nimport torch.profiler\nimport torch.utils.data\nimport torchvision.datasets\nimport torchvision.models\nimport torchvision.transforms as T"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then prepare the input data. For this tutorial, we use the CIFAR10 dataset.\nTransform it to the desired format and use DataLoader to load each batch.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "transform = T.Compose(\n    [T.Resize(224),\n     T.ToTensor(),\n     T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\ntrain_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\ntrain_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next, create Resnet model, loss function, and optimizer objects.\nTo run on GPU, move model and loss to GPU device.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda:0\")\nmodel = torchvision.models.resnet18(pretrained=True).cuda(device)\ncriterion = torch.nn.CrossEntropyLoss().cuda(device)\noptimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\nmodel.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the training step for each batch of input data.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def train(data):\n    inputs, labels = data[0].to(device=device), data[1].to(device=device)\n    outputs = model(inputs)\n    loss = criterion(outputs, labels)\n    optimizer.zero_grad()\n    loss.backward()\n    optimizer.step()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 2. Use profiler to record execution events\n\nThe profiler is enabled through the context manager and accepts several parameters,\nsome of the most useful are:\n\n- ``schedule`` - callable that takes step (int) as a single parameter\n  and returns the profiler action to perform at each step.\n\n  In this example with ``wait=1, warmup=1, active=3, repeat=2``,\n  profiler will skip the first step/iteration,\n  start warming up on the second,\n  record the following three iterations,\n  after which the trace will become available and on_trace_ready (when set) is called.\n  In total, the cycle repeats twice. Each cycle is called a \"span\" in TensorBoard plugin.\n\n  During ``wait`` steps, the profiler is disabled.\n  During ``warmup`` steps, the profiler starts tracing but the results are discarded.\n  This is for reducing the profiling overhead.\n  The overhead at the beginning of profiling is high and easy to bring skew to the profiling result.\n  During ``active`` steps, the profiler works and records events.\n- ``on_trace_ready`` - callable that is called at the end of each cycle;\n  In this example we use ``torch.profiler.tensorboard_trace_handler`` to generate result files for TensorBoard.\n  After profiling, result files will be saved into the ``./log/resnet18`` directory.\n  Specify this directory as a ``logdir`` parameter to analyze profile in TensorBoard.\n- ``record_shapes`` - whether to record shapes of the operator inputs.\n- ``profile_memory`` - Track tensor memory allocation/deallocation. Note, for old version of pytorch with version\n  before 1.10, if you suffer long profiling time, please disable it or upgrade to new version.\n- ``with_stack`` - Record source information (file and line number) for the ops.\n  If the TensorBoard is launched in VSCode ([reference](https://code.visualstudio.com/docs/datascience/pytorch-support#_tensorboard-integration)),\n  clicking a stack frame will navigate to the specific code line.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "with torch.profiler.profile(\n        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),\n        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),\n        record_shapes=True,\n        profile_memory=True,\n        with_stack=True\n) as prof:\n    for step, batch_data in enumerate(train_loader):\n        if step >= (1 + 1 + 3) * 2:\n            break\n        train(batch_data)\n        prof.step()  # Need to call this at the end of each step to notify profiler of steps' boundary."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Alternatively, the following non-context manager start/stop is supported as well.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "prof = torch.profiler.profile(\n        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),\n        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),\n        record_shapes=True,\n        with_stack=True)\nprof.start()\nfor step, batch_data in enumerate(train_loader):\n    if step >= (1 + 1 + 3) * 2:\n        break\n    train(batch_data)\n    prof.step()\nprof.stop()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 3. Run the profiler\n\nRun the above code. The profiling result will be saved under ``./log/resnet18`` directory.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 4. Use TensorBoard to view results and analyze model performance\n\nInstall PyTorch Profiler TensorBoard Plugin.\n\n::\n\n    pip install torch_tb_profiler\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Launch the TensorBoard.\n\n::\n\n    tensorboard --logdir=./log\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Open the TensorBoard profile URL in Google Chrome browser or Microsoft Edge browser.\n\n::\n\n    http://localhost:6006/#pytorch_profiler\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "You could see Profiler plugin page as shown below.\n\n- Overview\n<img src=\"file://../../_static/img/profiler_overview1.png\" scale=\"25 %\">\n\nThe overview shows a high-level summary of model performance.\n\nThe \"GPU Summary\" panel shows the GPU configuration, GPU usage and Tensor Cores usage.\nIn this example, the GPU Utilization is low.\nThe details of these metrics are [here](https://github.com/pytorch/kineto/blob/main/tb_plugin/docs/gpu_utilization.md).\n\nThe \"Step Time Breakdown\" shows distribution of time spent in each step over different categories of execution.\nIn this example, you can see the ``DataLoader`` overhead is significant.\n\nThe bottom \"Performance Recommendation\" uses the profiling data\nto automatically highlight likely bottlenecks,\nand gives you actionable optimization suggestions.\n\nYou can change the view page in left \"Views\" dropdown list.\n\n<img src=\"file://../../_static/img/profiler_views_list.png\"\n   :alt:>\n\n\n- Operator view\nThe operator view displays the performance of every PyTorch operator\nthat is executed either on the host or device.\n\n<img src=\"file://../../_static/img/profiler_operator_view.png\" scale=\"25 %\">\nThe \"Self\" duration does not include its child operators\u2019 time.\nThe \"Total\" duration includes its child operators\u2019 time.\n\n- View call stack\nClick the \"View Callstack\" of an operator, the operators with same name but different call stacks will be shown.\nThen click a \"View Callstack\" in this sub-table, the call stack frames will be shown.\n\n<img src=\"file://../../_static/img/profiler_callstack.png\" scale=\"25 %\">\n\nIf the TensorBoard is launched inside VSCode\n([Launch Guide](https://devblogs.microsoft.com/python/python-in-visual-studio-code-february-2021-release/#tensorboard-integration)),\nclicking a call stack frame will navigate to the specific code line.\n\n<img src=\"file://../../_static/img/profiler_vscode.png\" scale=\"25 %\">\n\n\n- Kernel view\nThe GPU kernel view shows all kernels\u2019 time spent on GPU.\n\n<img src=\"file://../../_static/img/profiler_kernel_view.png\" scale=\"25 %\">\nTensor Cores Used:\nWhether this kernel uses Tensor Cores.\n\nMean Blocks per SM:\nBlocks per SM = Blocks of this kernel / SM number of this GPU.\nIf this number is less than 1, it indicates the GPU multiprocessors are not fully utilized.\n\"Mean Blocks per SM\" is weighted average of all runs of this kernel name, using each run\u2019s duration as weight.\n\nMean Est. Achieved Occupancy:\nEst. Achieved Occupancy is defined in this column\u2019s tooltip.\nFor most cases such as memory bandwidth bounded kernels, the higher the better.\n\"Mean Est. Achieved Occupancy\" is weighted average of all runs of this kernel name,\nusing each run\u2019s duration as weight.\n\n- Trace view\nThe trace view shows timeline of profiled operators and GPU kernels.\nYou can select it to see details as below.\n\n<img src=\"file://../../_static/img/profiler_trace_view1.png\" scale=\"25 %\">\n\nYou can move the graph and zoom in/out with the help of right side toolbar.\nAnd keyboard can also be used to zoom and move around inside the timeline.\nThe \u2018w\u2019 and \u2018s\u2019 keys zoom in centered around the mouse,\nand the \u2018a\u2019 and \u2018d\u2019 keys move the timeline left and right.\nYou can hit these keys multiple times until you see a readable representation.\n\nIf a backward operator's \"Incoming Flow\" field is with value \"forward correspond to backward\",\nyou can click the text to get its launching forward operator.\n\n<img src=\"file://../../_static/img/profiler_trace_view_fwd_bwd.png\" scale=\"25 %\">\n\nIn this example, we can see the event prefixed with ``enumerate(DataLoader)`` costs a lot of time.\nAnd during most of this period, the GPU is idle.\nBecause this function is loading data and transforming data on host side,\nduring which the GPU resource is wasted.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 5. Improve performance with the help of profiler\n\nAt the bottom of \"Overview\" page, the suggestion in \"Performance Recommendation\" hints the bottleneck is DataLoader.\nThe PyTorch DataLoader uses single process by default.\nUser could enable multi-process data loading by setting the parameter ``num_workers``.\n[Here](https://pytorch.org/docs/stable/data.html#single-and-multi-process-data-loading) is more details.\n\nIn this example, we follow the \"Performance Recommendation\" and set ``num_workers`` as below,\npass a different name such as ``./log/resnet18_4workers`` to ``tensorboard_trace_handler``, and run it again.\n\n::\n\n    train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=4)\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then let\u2019s choose the recently profiled run in left \"Runs\" dropdown list.\n\n<img src=\"file://../../_static/img/profiler_overview2.png\" scale=\"25 %\">\n\nFrom the above view, we can find the step time is reduced to about 76ms comparing with previous run's 132ms,\nand the time reduction of ``DataLoader`` mainly contributes.\n\n<img src=\"file://../../_static/img/profiler_trace_view2.png\" scale=\"25 %\">\n\nFrom the above view, we can see that the runtime of ``enumerate(DataLoader)`` is reduced,\nand the GPU utilization is increased.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 6. Analyze performance with other advanced features\n\n- Memory view\nTo profile memory, ``profile_memory`` must be set to ``True`` in arguments of ``torch.profiler.profile``.\n\nYou can try it by using existing example on Azure\n\n::\n\n    pip install azure-storage-blob\n    tensorboard --logdir=https://torchtbprofiler.blob.core.windows.net/torchtbprofiler/demo/memory_demo_1_10\n\nThe profiler records all memory allocation/release events and allocator's internal state during profiling.\nThe memory view consists of three components as shown in the following.\n\n<img src=\"file://../../_static/img/profiler_memory_view.png\" scale=\"25 %\">\n\nThe components are memory curve graph, memory events table and memory statistics table, from top to bottom, respectively.\n\nThe memory type could be selected in \"Device\" selection box.\nFor example, \"GPU0\" means the following table only shows each operator's memory usage on GPU 0, not including CPU or other GPUs.\n\nThe memory curve shows the trends of memory consumption. The \"Allocated\" curve shows the total memory that is actually\nin use, e.g., tensors. In PyTorch, caching mechanism is employed in CUDA allocator and some other allocators. The\n\"Reserved\" curve shows the total memory that is reserved by the allocator. You can left click and drag on the graph\nto select events in the desired range:\n\n<img src=\"file://../../_static/img/profiler_memory_curve_selecting.png\" scale=\"25 %\">\n\nAfter selection, the three components will be updated for the restricted time range, so that you can gain more\ninformation about it. By repeating this process, you can zoom into a very fine-grained detail. Right click on the graph\nwill reset the graph to the initial state.\n\n<img src=\"file://../../_static/img/profiler_memory_curve_single.png\" scale=\"25 %\">\n\nIn the memory events table, the allocation and release events are paired into one entry. The \"operator\" column shows\nthe immediate ATen operator that is causing the allocation. Notice that in PyTorch, ATen operators commonly use\n``aten::empty`` to allocate memory. For example, ``aten::ones`` is implemented as ``aten::empty`` followed by an\n``aten::fill_``. Solely display the opeartor name as ``aten::empty`` is of little help. It will be shown as\n``aten::ones (aten::empty)`` in this special case. The \"Allocation Time\", \"Release Time\" and \"Duration\"\ncolumns' data might be missing if the event occurs outside of the time range. \n\nIn the memory statistics table, the \"Size Increase\" column sums up all allocation size and minus all the memory\nrelease size, that is, the net increase of memory usage after this operator. The \"Self Size Increase\" column is\nsimilar to \"Size Increase\", but it does not count children operators' allocation. With regards to ATen operators'\nimplementation detail, some operators might call other operators, so memory allocations can happen at any level of the\ncall stack. That says, \"Self Size Increase\" only count the memory usage increase at current level of call stack.\nFinally, the \"Allocation Size\" column sums up all allocation without considering the memory release.\n\n- Distributed view\nThe plugin now supports distributed view on profiling DDP with NCCL/GLOO as backend.\n\nYou can try it by using existing example on Azure:\n\n::\n\n    pip install azure-storage-blob\n    tensorboard --logdir=https://torchtbprofiler.blob.core.windows.net/torchtbprofiler/demo/distributed_bert\n\n<img src=\"file://../../_static/img/profiler_distributed_view.png\" scale=\"25 %\">\n\nThe \"Computation/Communication Overview\" shows computation/communication ratio and their overlapping degree.\nFrom this view, User can figure out load balance issue among workers.\nFor example, if the computation + overlapping time of one worker is much larger than others,\nthere may be a problem of load balance or this worker may be a straggler.\n\nThe \"Synchronizing/Communication Overview\" shows the efficiency of communication.\n\"Data Transfer Time\" is the time for actual data exchanging.\n\"Synchronizing Time\" is the time for waiting and synchronizing with other workers.\n\nIf one worker\u2019s \"Synchronizing Time\" is much shorter than that of other workers\u2019,\nthis worker may be a straggler which may have more computation workload than other workers\u2019.\n\nThe \"Communication Operations Stats\" summarizes the detailed statistics of all communication ops in each worker.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Learn More\n\nTake a look at the following documents to continue your learning,\nand feel free to open an issue [here](https://github.com/pytorch/kineto/issues).\n\n-  [Pytorch TensorBoard Profiler github](https://github.com/pytorch/kineto/tree/master/tb_plugin)\n-  [torch.profiler API](https://pytorch.org/docs/master/profiler.html)\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}